from selenium.webdriver.common.by import By

from base.base_page import BasePage
from utils.captcha_ocr import recognize_captcha


class LoginPage(BasePage):

    _USERNAME_INPUT = (By.ID, "userName")
    _PASSWORD_INPUT = (By.ID, "password")
    _VERIFY_INPUT = (By.ID, "randCode")
    _LOGIN_BUTTON = (By.ID, "but_login")
    _CAPTCHA_IMG = (By.ID, "randCodeImage")

    # 构造函数删掉 self.wait，直接用父类的
    def enter_username(self, username: str):
        self.send_keys(self._USERNAME_INPUT, username)
        return self

    def enter_password(self, password: str):
        self.send_keys(self._PASSWORD_INPUT, password)
        return self

    def enter_verify_code(self, code: str):
        self.send_keys(self._VERIFY_INPUT, code)
        return self

    def click_login(self):
        self.click(self._LOGIN_BUTTON)

    def get_verify_code(self) -> str:
        img_bytes = self.find(self._CAPTCHA_IMG).screenshot_as_png
        return recognize_captcha(img_bytes)

    def login(self, username: str, password: str, verify_code: str):
        self.enter_username(username) \
            .enter_password(password) \
            .enter_verify_code(verify_code) \
            .click_login()
